#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/4/7 10:58
# @Author  : CHEN Wang
# @Site    :
# @File    : base_series.py
# @Software: PyCharm

"""
脚本说明：
Series 类
这个是下面以DataFrame为输入的基础函数
return pd.Series format
"""

from functools import reduce
import math
import numpy as np
import pandas as pd

try:
    import talib
except:
    print('PLEASE install TALIB to call these methods')
    pass


def MA(Series, N):
    return pd.Series.rolling(Series, N).mean()


def DIFF(Series, N=1):
    return pd.Series(Series).diff(N)


def HHV(Series, N):
    """
    过去 N 天 X 的最大值，比如 HHV(HIGH,N)为过去 N 天的最高价
    :param Series:
    :param N:
    :return:
    """
    return pd.Series(Series).rolling(N).max()


def LLV(Series, N):
    return pd.Series(Series).rolling(N).min()


def SUM(Series, N):
    return pd.Series.rolling(Series, N).sum()


def CUMSUM(Series, N=1):
    return pd.Series.expanding(Series, min_periods=N).sum()


def CUMPROD(Series):
    return Series.cumprod()


def ABS(Series):
    return abs(Series)


def ROLLMAX(Series, N):
    return pd.Series.rolling(Series, N).max()


def MAX(A, B):
    var = IF(A > B, A, B)
    return var


def ROLLMIN(Series, N):
    return pd.Series.rolling(Series, N).min()


def MIN(A, B):
    var = IF(A < B, A, B)
    return var


def SINGLE_CROSS(A, B):
    if A.iloc[-2] < B.iloc[-2] and A.iloc[-1] > B.iloc[-1]:
        return True
    else:
        return False


def EXIST(Series, N=5):
    '''
    n日内是否存在某一值，
    输入值为true or false的Series
    '''
    res = pd.DataFrame(Series) + 0
    res = res.rolling(N).sum() > 0
    res = res[res.columns[0]]
    return res


def EVERY(Series, N=5):
    '''
    n日内是否一直存在某一值，
    输入值为true or false的Series
    '''
    res = pd.DataFrame(Series) + 0
    res = res.rolling(N).sum() > N - 1
    res = res[res.columns[0]]
    return res


def CROSS(A, B):
    """A<B then A>B  A上穿B B下穿A
    Arguments:
        A {[type]} -- [description]
        B {[type]} -- [description]
    Returns:
        [type] -- [description]
    """

    var = np.where(A < B, 1, 0)
    try:
        index = A.index
    except:
        index = B.index
    return (pd.Series(var, index=index).diff() < 0).apply(int)


def CROSS_STATUS(A, B):
    """
    A 穿过 B 产生持续的 1 序列信号
    """
    return np.where(A > B, 1, 0)


def FILTER(COND, N):
    k1 = pd.Series(np.where(COND, 1, 0), index=COND.index)
    idx = k1[k1 == 1].index.codes[0]
    needfilter = pd.Series(idx, index=idx)
    afterfilter = needfilter.diff().apply(lambda x: False if x < N else True)
    k1.iloc[afterfilter[afterfilter].index] = 2
    return k1.apply(lambda x: 1 if x == 2 else 0)


def COUNT(COND, N):
    """
    2018/05/23 修改
    参考https://github.com/QUANTAXIS/QUANTAXIS/issues/429
    现在返回的是series
    """
    return pd.Series(np.where(COND, 1, 0), index=COND.index).rolling(N).sum()


def IF(COND, V1, V2):
    """
    如果满足条件 COND 则取值为 A，否则为 B
    :param COND:
    :param V1:
    :param V2:
    :return:
    """
    var = np.where(COND, V1, V2)
    try:
        try:
            index = V1.index
        except:
            index = COND.index
    except:
        index = V2.index
    return pd.Series(var, index=index)


def IFLAST(COND, V1):
    """
    如果满足条件 COND 则取值为 V1，否则为res的前一个值
    :param COND:
    :param V1:
    :return:
    """
    res = V1.copy()
    for i, idx in enumerate(res.index):
        if i == 0:
            res[idx] = V1[idx]
        else:
            if COND[idx]:
                res[idx] = V1[idx]
            else:
                res[idx] = res.iloc[i-1]

    return res


def IFAND(COND1, COND2, V1, V2):
    var = np.where(np.logical_and(COND1, COND2), V1, V2)
    return pd.Series(var, index=V1.index)


def IFOR(COND1, COND2, V1, V2):
    var = np.where(np.logical_or(COND1, COND2), V1, V2)
    return pd.Series(var, index=V1.index)


def REF(Series, N):
    """
    N 天前 X 的值，比如 REF(CLOSE,4)为 4 天前的收盘价

    :param Series:
    :param N:
    :return:
    """
    return Series.shift(N)


def LAST(COND, N1, N2):
    """表达持续性
    从前N1日到前N2日一直满足COND条件
    Arguments:
        COND {[type]} -- [description]
        N1 {[type]} -- [description]
        N2 {[type]} -- [description]
    """
    N2 = 1 if N2 == 0 else N2
    assert N2 > 0
    assert N1 > N2
    return COND.iloc[-N1:-N2].all()


def STD(Series, N):
    return pd.Series.rolling(Series, N).std()


def AVEDEV(Series, N):
    """
    平均绝对偏差 mean absolute deviation
    修正: 2018-05-25
    之前用mad的计算模式依然返回的是单值
    """
    return Series.rolling(N).apply(lambda x: (np.abs(x - x.mean())).mean(), raw=True)


def MACD(Series, FAST, SLOW, MID):
    """macd指标 仅适用于Series
    对于DATAFRAME的应用请使用QA_indicator_macd
    """
    EMAFAST = EMA(Series, FAST)
    EMASLOW = EMA(Series, SLOW)
    DIFF = EMAFAST - EMASLOW
    DEA = EMA(DIFF, MID)
    MACD = (DIFF - DEA) * 2
    DICT = {'DIFF': DIFF, 'DEA': DEA, 'MACD': MACD}
    VAR = pd.DataFrame(DICT)
    return VAR


def BBIBOLL(Series, N1, N2, N3, N4, N, M):  # 多空布林线

    bbiboll = BBI(Series, N1, N2, N3, N4)
    UPER = bbiboll + M * STD(bbiboll, N)
    DOWN = bbiboll - M * STD(bbiboll, N)
    DICT = {'BBIBOLL': bbiboll, 'UPER': UPER, 'DOWN': DOWN}
    VAR = pd.DataFrame(DICT)
    return VAR


def BBI(Series, N1, N2, N3, N4):
    '多空指标'

    bbi = (MA(Series, N1) + MA(Series, N2) +
           MA(Series, N3) + MA(Series, N4)) / 4
    DICT = {'BBI': bbi}
    VAR = pd.DataFrame(DICT)
    return VAR


def BARLAST(cond, yes=True):
    """支持MultiIndex的cond和DateTimeIndex的cond
    条件成立  yes= True 或者 yes=1 根据不同的指标自己定
    最后一次条件成立  到 当前到周期数
    Arguments:
        cond {[type]} -- [description]
    """
    if isinstance(cond.index, pd.MultiIndex):
        return len(cond) - cond.index.levels[0].tolist().index(cond[cond == yes].index[-1][0]) - 1
    elif isinstance(cond.index, pd.DatetimeIndex):
        return len(cond) - cond.index.tolist().index(cond[cond == yes].index[-1]) - 1


def BARLAST_EXIST(cond, yes=True):
    """
    上一次条件成立   持续到当前到数量
    支持MultiIndex的cond和DateTimeIndex的cond
    条件成立  yes= True 或者 yes=1 根据不同的指标自己定
    Arguments:
        cond {[type]} -- [description]
    """
    if isinstance(cond.index, pd.MultiIndex):
        return len(cond) - cond.index.levels[0].tolist().index(cond[cond != yes].index[-1][0]) - 1
    elif isinstance(cond.index, pd.DatetimeIndex):
        return len(cond) - cond.index.tolist().index(cond[cond != yes].index[-1]) - 1


def XARROUND(x, y): return np.round(
    y * (round(x / y - math.floor(x / y) + 0.00000000001) + math.floor(x / y)), 2)


def RENKO(Series, N, condensed=True):
    last_price = Series[0]
    chart = [last_price]
    for price in Series:
        bricks = math.floor(abs(price - last_price) / N)
        if bricks == 0:
            if condensed:
                chart.append(chart[-1])
            continue
        sign = int(np.sign(price - last_price))
        chart += [sign * (last_price + (sign * N * x)) for x in range(1, bricks + 1)]
        last_price = abs(chart[-1])

    return pd.Series(chart)


def RENKOP(Series, N, condensed=True):
    last_price = Series[0]
    chart = [last_price]
    for price in Series:
        inc = (price - last_price) / last_price
        # print(inc)
        if abs(inc) < N:
            # if condensed:
            #     chart.append(chart[-1])
            continue
        sign = int(np.sign(price - last_price))
        bricks = math.floor(inc / N)
        # print(bricks)
        # print((N * (price-last_price)) / inc)
        step = math.floor((N * (price - last_price)) / inc)
        print(step)
        # print(sign)
        chart += [sign * (last_price + (sign * step * x))
                  for x in range(1, abs(bricks) + 1)]
        last_price = abs(chart[-1])
    return pd.Series(chart)


def CMO(Series, timeperiod=14):
    res = talib.CMO(Series.values, timeperiod)
    return pd.Series(res, index=Series.index)


def BBANDS(Series, timeperiod=5, nbdevup=2, nbdevdn=2, matype=0):
    up, middle, low = talib.BBANDS(
        Series.values, timeperiod, nbdevup, nbdevdn, matype)
    return pd.Series(up, index=Series.index), pd.Series(middle, index=Series.index), pd.Series(low, index=Series.index)


def BETA(SeriesA, SeriesB, timeperiod=5):
    res = talib.BETA(SeriesA.values, SeriesB.values, timeperiod)  # Beta
    return pd.Series(res, index=SeriesA.index)


def CORREL(SeriesA, SeriesB, timeperiod=5):
    res = talib.BETA(SeriesA.values, SeriesB.values, timeperiod)  # Pearson's Correlation Coefficient (r)
    return pd.Series(res, index=SeriesA.index)


def HT_DCPERIOD(Series):
    res = talib.HT_DCPERIOD(Series.values)
    return pd.Series(res, index=Series.index)


def HT_DCPHASE(Series):
    res = talib.HT_DCPHASE(Series.values)
    return pd.Series(res, index=Series.index)


def HT_PHASOR(Series):
    res = talib.HT_PHASOR(Series.values)
    return pd.Series(res, index=Series.index)


def HT_SINE(Series):
    res = talib.HT_SINE(Series.values)
    return pd.Series(res, index=Series.index)


def HT_TRENDLINE(Series):
    res = talib.HT_TRENDLINE(Series.values)
    return pd.Series(res, index=Series.index)


def HT_TRENDMODE(Series):
    res = talib.HT_TRENDMODE(Series.values)
    return pd.Series(res, index=Series.index)


def KAMA(Series, timeperiod=30):
    res = talib.KAMA(Series.values, timeperiod)
    return pd.Series(res, index=Series.index)


def LINEARREG(Series, timeperiod=14):
    res = talib.LINEARREG(Series.values, timeperiod)  # Linear Regression
    return pd.Series(res, index=Series.index)


def LINEARREG_ANGLE(Series, timeperiod=14):
    res = talib.LINEARREG_ANGLE(Series.values, timeperiod)  # Linear Regression Angle
    return pd.Series(res, index=Series.index)


def LINEARREG_INTERCEPT(Series, timeperiod=14):
    res = talib.LINEARREG_INTERCEPT(Series.values, timeperiod)  # Linear Regression Intercept
    return pd.Series(res, index=Series.index)


def LINEARREG_SLOPE(Series, timeperiod=14):
    res = talib.LINEARREG_SLOPE(Series.values, timeperiod)  # Linear Regression Slope
    return pd.Series(res, index=Series.index)


def MACDEXT(Series, fastperiod=12, fastmatype=0, slowperiod=26, slowmatype=0, signalperiod=9, signalmatype=0):
    macd, macdsignal, macdhist = talib.MACDEXT(
        Series.values, fastperiod, fastmatype, slowperiod, slowmatype, signalperiod, signalmatype)
    return pd.Series(macd, index=Series.index), pd.Series(macdsignal, index=Series.index), pd.Series(macdhist,
                                                                                                     index=Series.index)


def MACDFIX(Series, timeperiod=9):
    macd, macdsignal, macdhist = talib.MACDFIX(Series.values, timeperiod)
    return pd.Series(macd, index=Series.index), pd.Series(macdsignal, index=Series.index), pd.Series(macdhist,
                                                                                                     index=Series.index)


def MAMA(Series, fastlimit=0.5, slowlimit=0.05):
    mama, fama = talib.MAMA(Series.values, fastlimit, slowlimit)
    return pd.Series(mama, index=Series.index), pd.Series(fama, index=Series.index)


def SMA(Series, N, M=1):
    """
    SMA = Cn*M/N+(1-M/N)*SMAn-1
    TA-Lib的SMA同国内量化交易普遍使用的SMA是不同的, 而是等同MA
    威廉SMA算法  参考https://www.joinquant.com/post/867

    特殊情况：
    M = 1 时， 等同于MA
    M = 2 时， 等同于EMA

    """
    ret = []
    i = 1
    length = len(Series)
    # 跳过X中前面几个 nan 值
    while i < length:
        if np.isnan(Series.iloc[i]):
            i += 1
        else:
            break
    preY = Series.iloc[i]  # Y'
    ret.append(preY)
    while i < length:
        Y = (M * Series.iloc[i] + (N - M) * preY) / float(N)
        ret.append(Y)
        preY = Y
        i += 1
    return pd.Series(ret, index=Series.tail(len(ret)).index)


def EMA(Series, timeperiod=30):
    res = talib.EMA(Series.values, timeperiod)
    return pd.Series(res, index=Series.index)
    # return pd.Series.ewm(Series, span=N, min_periods=N - 1, adjust=True).mean()


def DEMA(Series, timeperiod=30):
    res = talib.DEMA(Series.values, timeperiod)
    return pd.Series(res, index=Series.index)


def STDDEV(Series, timeperiod=5, nbdev=1):
    return pd.Series(talib.STDDEV(Series.values, timeperiod, nbdev), index=Series.index)  # Standard Deviation


def VAR(Series, timeperiod=5, nbdev=1):
    return pd.Series(talib.VAR(Series.values, timeperiod, nbdev), index=Series.index)  # Variance


def TSF(Series, timeperiod=5, nbdev=1):
    return pd.Series(talib.TSF(Series.values, timeperiod, nbdev), index=Series.index)  # Time Series Forecast


def STOCHRSI(Series, timeperiod=14, fastk_period=5, fastd_period=3, fastd_matype=0):
    fastk, fastd = talib.STOCHRSI(
        Series.values, fastk_period, fastd_period, fastd_matype)
    return pd.Series(fastk, index=Series.index), pd.Series(fastd, index=Series.index)


if __name__ == '__main__':
    import os
    from quant_researcher.quant.project_tool.localize import DATA_DIR
    file_name = os.path.join(DATA_DIR, f'BTC_history_ohlcvm')
    all_data = pd.read_excel(f'{file_name}.xlsx', index_col='end_date')

    open = all_data.open
    close = all_data.close
    high = all_data.high
    low = all_data.low
    periods = 3